package org.dromara.easyai.extensions.cuda;

import org.dromara.easyai.matrixTools.Matrix;

/**
 * @author: wen.y
 * @date: 2025/1/7
 */
public class MatrixUtil {

	/**
	 * 将矩阵转换为行主序数组
	 * @param matrix
	 * @return
	 * @throws Exception
	 */
	public static float[] toRowMajorArray(Matrix matrix) throws Exception {
		float[] res = new float[matrix.getX() * matrix.getY()];
		for (int i = 0; i < matrix.getX(); i++) {
			for (int j = 0; j < matrix.getY(); j++) {
				res[i * matrix.getY() + j] = matrix.getNumber(i, j);
			}
		}
		return res;
	}


	/**
	 * 将矩阵转换为列主序数组
	 * @param matrix
	 * @return
	 */
	public static float[] toColumnMajorArray(Matrix matrix) throws Exception {
		float[] res = new float[matrix.getX() * matrix.getY()];
		for (int i = 0; i < matrix.getX(); i++) {
			for (int j = 0; j < matrix.getY(); j++) {
				res[j * matrix.getX() + i] = matrix.getNumber(i, j);
			}
		}
		return res;
	}

	/**
	 * 生成一个n阶单位矩阵（列主序）
	 * @param n
	 * @return
	 */
	public static double[] makeFp64UnitColumnMajorArray(int n) {
		double[] res = new double[n * n];
		for (int i = 0; i < n; i++) {
			for (int j = 0; j < n; j++) {
				res[j * n + i] = (i == j) ? 1: 0;
			}
		}
		return res;
	}

	/**
	 * 生成一个n阶单位矩阵（列主序）
	 * @param n
	 * @return
	 */
	public static float[] makeFp32UnitColumnMajorArray(int n) {
		float[] res = new float[n * n];
		for (int i = 0; i < n; i++) {
			for (int j = 0; j < n; j++) {
				res[j * n + i] = (i == j) ? 1: 0;
			}
		}
		return res;
	}

	public static double[] makeFp64Array(int rows, int cols, double value) {
		double[] res = new double[rows * cols];
		for (int i = 0; i < rows; i++) {
			for (int j = 0; j < cols; j++) {
				res[j * rows + i] = value;
			}
		}
		return res;
	}

	public static float[] makeFp32Array(int rows, int cols, float value) {
		float[] res = new float[rows * cols];
		for (int i = 0; i < rows; i++) {
			for (int j = 0; j < cols; j++) {
				res[j * rows + i] = value;
			}
		}
		return res;
	}

	/**
	 * 将列主序数组转换为矩阵
	 * @param array
	 * @param rows
	 * @param cols
	 * @return
	 * @throws Exception
	 */
	public static Matrix columnMajorArrayToMatrix(float[] array, int rows, int cols) throws Exception {
		Matrix matrix = new Matrix(rows, cols);
		for (int i = 0; i < rows; i++) {
			for (int j = 0; j < cols; j++) {
				matrix.setNub(i, j, array[j * rows + i]);
			}
		}
		return matrix;
	}

	/**
	 * 将行主序数组转换为矩阵
	 * @param array
	 * @param rows
	 * @param cols
	 * @return
	 * @throws Exception
	 */
	public static Matrix rowMajorArrayToMatrix(float[] array, int rows, int cols) throws Exception {
		Matrix matrix = new Matrix(rows, cols);
		for (int i = 0; i < rows; i++) {
			for (int j = 0; j < cols; j++) {
				matrix.setNub(i, j, array[i * cols + j]);
			}
		}
		return matrix;
	}

	public static void setRowMajorArray(Matrix matrix, float[] array) throws Exception {
		for (int i = 0; i < matrix.getX(); i++) {
			for (int j = 0; j < matrix.getY(); j++) {
				matrix.setNub(i, j, array[i * matrix.getY() + j]);
			}
		}
	}

	public static void setColumnMajorArray(Matrix matrix, float[] array) throws Exception {
		for (int i = 0; i < matrix.getX(); i++) {
			for (int j = 0; j < matrix.getY(); j++) {
				matrix.setNub(i, j, array[j * matrix.getX() + i]);
			}
		}
	}

}
